import numpy as np
from osgeo import gdal, gdalconst


#  ===============================================================
#  计算生物多样性
# 生物多样性是由区域生境质量、植被净初生产力和景观结构指数的加权函数式组成的，其概念框架为：
#            BIx=Qx;×β1+NPPx×β₂+Sx×β₃
# 式中，BI表示栅格单元x的生物多样性指数，Qx;为第j类土地利用/覆被类型(地类景观)在第x个栅格的生境质量，
# NPP为第x个栅格的植被净初生产力年值，Sx为景观尺度上第x个栅格的景观结构指数，βx为权重系数，通过结合研
# 究区实地特点，利用层次分析法对生境质量、植被净初生产力和景观结构指数权重值分别赋系数，默认为0.5、0.3
# 和0.2,可自行更改。各指标均进行标准化处理。
# 输出数据为2.2.5生境质量生成的Qx;生境质量栅格，2.2.1生成的NPP生产力栅
# 格和2.1.2生成的LS;景观结构指数栅格，输出数据为BIx生物多样性栅格数据。
#  ===============================================================

def resample_to_memory(input_path, target_width, target_height, resample_alg="bilinear"):
    """
    将栅格图像重采样到目标尺寸，并返回内存中的栅格数据。

    参数:
    input_path (str): 输入栅格文件的路径。
    target_width (int): 目标宽度（列数）。
    target_height (int): 目标高度（行数）。
    resample_alg (str): 重采样方法，可选值为 "near"（最近邻）、"bilinear"（双线性）、"cubic"（立方卷积）等。

    返回:
    numpy.ndarray: 重采样后的栅格数据。
    """
    # 打开输入栅格文件
    input_dataset = gdal.Open(input_path, gdalconst.GA_ReadOnly)
    if input_dataset is None:
        raise FileNotFoundError(f"无法打开文件: {input_path}")

    # 获取输入栅格的地理变换和投影
    input_geotransform = input_dataset.GetGeoTransform()
    input_projection = input_dataset.GetProjection()

    # 计算新的地理变换
    # 地理变换公式: (左上角X坐标, X方向分辨率, 旋转参数, 左上角Y坐标, 旋转参数, Y方向分辨率)
    x_res = input_geotransform[1] * (input_dataset.RasterXSize / target_width)  # 调整X方向分辨率
    y_res = input_geotransform[5] * (input_dataset.RasterYSize / target_height)  # 调整Y方向分辨率
    new_geotransform = (
        input_geotransform[0],  # 左上角X坐标不变
        x_res,                  # 新的X方向分辨率
        input_geotransform[2],  # 旋转参数不变
        input_geotransform[3],  # 左上角Y坐标不变
        input_geotransform[4],  # 旋转参数不变
        y_res                   # 新的Y方向分辨率
    )

    # 创建内存中的输出数据集
    driver = gdal.GetDriverByName("MEM")
    output_dataset = driver.Create("", target_width, target_height, 1, gdalconst.GDT_Float32)

    # 设置输出数据集的地理变换和投影
    output_dataset.SetGeoTransform(new_geotransform)
    output_dataset.SetProjection(input_projection)

    # 使用 gdal.Warp 进行重采样
    gdal.Warp(
        output_dataset,
        input_dataset,
        width=target_width,
        height=target_height,
        resampleAlg=resample_alg
    )

    # 读取重采样后的数据
    band = output_dataset.GetRasterBand(1)
    data = band.ReadAsArray()

    # 关闭数据集
    input_dataset = None
    output_dataset = None

    return data


def normalize_raster(data, nodata_value):
    """
    对栅格数据进行归一化，缩放到 [0, 1] 范围，忽略 NoData 值。

    参数:
    data (numpy.ndarray): 输入的栅格数据。
    nodata_value (float): NoData 值。

    返回:
    numpy.ndarray: 归一化后的栅格数据。
    """
    # 创建一个掩码，忽略 NoData 值
    masked_data = np.ma.masked_values(data, nodata_value)

    # 计算忽略 NoData 值后的最小值和最大值
    min_value = masked_data.min()
    max_value = masked_data.max()

    # 如果最大值等于最小值，返回全零数组
    if max_value == min_value:
        return np.zeros_like(data)

    # 归一化处理
    normalized_data = (masked_data - min_value) / (max_value - min_value)

    # 将掩码数组转换回普通数组，并将 NoData 值填充为原值
    result = normalized_data.filled(nodata_value)
    return result


def weighted_sum_of_rasters(raster_paths, weights, normalize=False):
    """
    读取三个GeoTIFF栅格图像，并将它们与给定的权重相乘后求和。

    参数:
    raster_paths (list): 三个GeoTIFF文件的路径列表。
    weights (list): 三个权重值的列表。
    normalize (bool): 是否对输入栅格数据进行归一化。

    返回:
    numpy.ndarray: 加权求和后的栅格图像。
    numpy.ndarray: 所有输入栅格的统一 NoData 掩码。
    """
    if len(raster_paths) != 3 or len(weights) != 3:
        raise ValueError("必须提供三个栅格路径和三个权重值。")

    # 初始化一个空列表来存储读取的栅格数据
    raster_data = []
    nodata_masks = []

    # 读取参考栅格的尺寸
    reference_dataset = gdal.Open(raster_paths[0], gdalconst.GA_ReadOnly)
    if reference_dataset is None:
        raise FileNotFoundError(f"无法打开文件: {raster_paths[0]}")
    target_width = reference_dataset.RasterXSize
    target_height = reference_dataset.RasterYSize
    reference_dataset = None

    # 读取每个栅格图像
    for path in raster_paths:
        dataset = gdal.Open(path, gdalconst.GA_ReadOnly)
        if dataset is None:
            raise FileNotFoundError(f"无法打开文件: {path}")

        # 检查栅格尺寸是否与参考栅格一致
        if dataset.RasterXSize != target_width or dataset.RasterYSize != target_height:
            # 如果不一致，重采样到参考栅格的尺寸
            print(f"栅格 {path} 的尺寸不一致，正在重采样...")
            data = resample_to_memory(path, target_width, target_height)
        else:
            # 如果尺寸一致，直接读取数据
            band = dataset.GetRasterBand(1)
            data = band.ReadAsArray()

        # 获取 NoData 值
        band = dataset.GetRasterBand(1)
        nodata_value = band.GetNoDataValue()
        if nodata_value is None:
            nodata_value = np.nan  # 如果没有 NoData 值，使用 NaN

        # 生成 NoData 掩码
        if nodata_value is not None:
            nodata_mask = (data == nodata_value)
        else:
            nodata_mask = np.zeros_like(data, dtype=bool)  # 如果没有 NoData 值，不进行掩码
        nodata_masks.append(nodata_mask)

        if normalize:
            data = normalize_raster(data, nodata_value)  # 归一化处理，考虑 NoData 值

        raster_data.append(data)

        # 关闭数据集
        dataset = None

    # 合并所有 NoData 掩码
    combined_nodata_mask = np.logical_or.reduce(nodata_masks)

    # 将权重与栅格数据相乘并求和
    weighted_sum = np.zeros_like(raster_data[0], dtype=np.float32)
    for i in range(3):
        weighted_sum += raster_data[i] * weights[i]

    return weighted_sum, combined_nodata_mask


def save_raster(output_path, data, reference_raster_path, nodata_value=None, nodata_mask=None):
    """
    将结果保存为新的GeoTIFF文件，并使用 NoData 掩码裁剪栅格。

    参数:
    output_path (str): 输出文件的路径。
    data (numpy.ndarray): 要保存的栅格数据。
    reference_raster_path (str): 参考栅格文件的路径（用于获取地理信息）。
    nodata_value (float): 输出文件的 NoData 值。
    nodata_mask (numpy.ndarray): NoData 掩码（可选）。
    """
    # 打开参考栅格文件
    reference_dataset = gdal.Open(reference_raster_path, gdalconst.GA_ReadOnly)
    if reference_dataset is None:
        raise FileNotFoundError(f"无法打开参考文件: {reference_raster_path}")

    # 如果提供了 NoData 掩码，将掩码区域设置为 NoData 值
    if nodata_mask is not None and nodata_value is not None:
        data[nodata_mask] = nodata_value

    # 获取参考文件的地理信息和元数据
    driver = gdal.GetDriverByName("GTiff")
    rows, cols = data.shape
    out_dataset = driver.Create(output_path, cols, rows, 1, gdalconst.GDT_Float32)

    # 设置地理变换和投影
    out_dataset.SetGeoTransform(reference_dataset.GetGeoTransform())
    out_dataset.SetProjection(reference_dataset.GetProjection())

    # 写入数据
    out_band = out_dataset.GetRasterBand(1)
    out_band.WriteArray(data)

    # 设置输出文件的 NoData 值
    if nodata_value is not None:
        out_band.SetNoDataValue(nodata_value)

    # 关闭数据集
    out_band.FlushCache()
    out_dataset = None
    reference_dataset = None

# 示例使用
if __name__ == "__main__":
    # =============================2.计算生物多样性指数===============================
    # 定义输入栅格文件路径和权重
    raster_paths = ["F:/test/多样性/NDVI2022.tif", "F:/test/多样性/NPP2022.tif", "F:/test/多样性/shdi2022.tif"]
    weights = [0.5, 0.3, 0.2]
    output_path = "F:/test/多样性/weighted_sum_result.tif"        # 保存结果到新的GeoTIFF文件

    # 是否对输入栅格数据进行归一化
    normalize_input = True  # 设置为 True 启用归一化，False 禁用归一化

    # 调用加权求和函数
    try:
        result, combined_nodata_mask = weighted_sum_of_rasters(raster_paths, weights, normalize=normalize_input)
        print("加权求和计算完成！")

        # 保存结果到新的GeoTIFF文件
        save_raster(output_path, result, raster_paths[0], nodata_value=np.nan, nodata_mask=combined_nodata_mask)
        print(f"结果已保存到: {output_path}")

    except Exception as e:
        print(f"发生错误: {e}")